import os

import pytest

from storage.kvstore.lmdb import LMDBKVStore


def test_get_put_lmdb(kv_path):
    store = LMDBKVStore(kv_path)
    with store.tx():
        store.put(b'key', b'value')
    with store.tx():
        assert store.get(b'key') == b'value'


def test_getitem_setitem_lmdb(kv_path):
    store = LMDBKVStore(kv_path)
    with store.tx():
        store[b'key'] = b'value'
    with store.tx():
        assert store[b'key'] == b'value'


def test_delete_lmdb(kv_path):
    store = LMDBKVStore(kv_path)
    with store.tx():
        store.put(b'key', b'value')

    with store.tx():
        del store[b'key']

    with store.tx():
        assert store.get(b'key') is None


@pytest.mark.benchmark
def test_benchmark_read_lmdb(benchmark, kv_path, get_keys, insert_keys):
    db = LMDBKVStore(str(kv_path))
    for k in insert_keys:
        db[k] = os.urandom(8)

    def benchmark_func():
        nonlocal db
        for k in get_keys:
            _ = db.get(k)

    benchmark(benchmark_func)


@pytest.mark.benchmark
def test_benchmark_write_lmdb(benchmark, kv_path, insert_keys):
    def benchmark_func():
        db = LMDBKVStore(str(kv_path))
        with db:
            for start in range(len(insert_keys) // 1000):
                keys = insert_keys[start * 1000:(start + 1) * 1000]
                data = {k: b'012345678' for k in keys}
                db.update(data)

    benchmark(benchmark_func)
